import sys
import os
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"
import numpy as np
import cv2 as cv
from PIL import Image

from atlas_utils.acl_resource import AclResource
from constants import *
from acl_model import Model

home_path = !echo ${HOME}
sys.path.append(os.path.join(home_path[0] , "jupyter-notebook/"))
print('System init success.')

acl_resource = AclResource()

acl_resource.init()


def init(self):
    """
    Init resource
    """
    print("init resource stage:")
    ret = acl.init()
    utils.check_ret("acl.init", ret)

  
    ret = acl.rt.set_device(self.device_id)
    utils.check_ret("acl.rt.set_device", ret)
    print("Set device n success.")

    #显式创建一个Context
    self.context, ret = acl.rt.create_context(self.device_id)
    utils.check_ret("acl.rt.create_context", ret)

    #创建一个Stream
    self.stream, ret = acl.rt.create_stream()
    utils.check_ret("acl.rt.create_stream", ret)

    #获取当前昇腾AI软件栈的运行模式
    self.run_mode, ret = acl.rt.get_run_mode()
    utils.check_ret("acl.rt.get_run_mode", ret)

    print("Init resource success")

def __del__(self):
    print("acl resource release all resource")
    resource_list.destroy()

    # 调用acl.rt.destroy_stream接口释放Stream
    if self.stream:
        acl.rt.destroy_stream(self.stream)
        print("acl resource release stream")

    # 调用acl.rt.destroy_context接口释放Context
    if self.context:
        acl.rt.destroy_context(self.context)
        print("acl resource release context")

    # 调用acl.rt.destroy_context接口释放Context
    acl.rt.reset_device(self.device_id)
    acl.finalize()
    print("Release acl resource success")

# 从文件加载离线模型数据
model_path = 'deeplabv3_plus.om'
model = Model(model_path)
model=model.npu()

MODEL_WIDTH = 512
MODEL_HEIGHT = 512

#图片路径
image_dir = '1.jpg'

IMG_EXT = ['.jpg', '.JPG', '.png', '.PNG', '.bmp', '.BMP', '.jpeg', '.JPEG']
images_list = [os.path.join(image_dir,img)for img in os.listdir(image_dir)if os.path.splitext(img)[1] in IMG_EXT]
images_list
Image.open(images_list[0])
type(images_list[0])


# 1. 读取图片
bgr_img = cv.imread(images_list[0])

# 2. 得到图片shape
orig_shape = bgr_img.shape[:2]
print("Shape of original image：",bgr_img.shape)

# 3.对图片进行resize，缩放至模型要求的宽高：
img = cv.resize(bgr_img, (MODEL_WIDTH, MODEL_HEIGHT))
print("The size of the model to be entered：",img.shape)

# 保存并展示resize之后的图片
if not os.path.isdir('./outputs'):
   os.mkdir('./outputs')
output_path = os.path.join("./outputs", "out_" + os.path.basename(images_list[0]))
cv.imwrite(output_path, img)
Image.open(output_path)
img = img.astype(np.int8).npu()
if not img.flags['C_CONTIGUOUS']:
        img = np.ascontiguousarray(img).npu()
img.shape
img.nbytes
# 根据构建好的模型输入数据进行模型推理
result_list = model.execute([img])
print(result_list)
# 1. 得到推理结果并且reshape成513*513，并且转换成uint8类型
result_img = result_list[0].reshape(513, 513)
result_img = result_img.astype('uint8')
result_img.shape

cv.imwrite(output_path, result_img)
Image.open(output_path)

# 2. 多通道图像融合
img = cv.merge((result_img, result_img, result_img))
print("Multi channel merged shape：",img.shape)

# 3. 将融合后的图像resize成原始图片的大小
bgr_img = cv.resize(img, (orig_shape[1], orig_shape[0]))
print("Shape of the original image size that is resized:",bgr_img.shape)
bgr_img = (bgr_img * 255)

# 保存图片
cv.imwrite(output_path, bgr_img)
Image.open(output_path)




